// Copyright 2022 jeff.li. and/or its affiliates.

#include <tbir/runtime/algorithm/prefix_matcher.h>

#include <memory>
#include <set>
#include <string>

#include <tbir/runtime/logging.h>
#include <tbir/runtime/utf8_util.h>

namespace tbir::runtime {

    PrefixMatcher::PrefixMatcher(const std::set<std::string> &dic) {
        if (dic.empty())
            return;
        std::vector<const char *> key;
        key.reserve(dic.size());
        for (const auto &it : dic) {
            key.push_back(it.data());
        }
        trie_ = std::make_unique<cedar_t>();
        MXCHECK_EQ(0, trie_->build(key.size(), const_cast<const char **>(&key[0]), nullptr, nullptr));
    }

    int PrefixMatcher::PrefixMatch(const char *w, size_t w_len, bool *found) const {
        if (trie_ == nullptr) {
            if (found) {
                *found = false;
            }
            return std::min<int>(w_len, OneCharLen(w));
        }
        constexpr int kResultSize = 64;
        cedar_t::result_pair_type trie_results[kResultSize];
        const int num_nodes = trie_->commonPrefixSearch(w, trie_results, kResultSize, w_len);

        if (found) {
            *found = (num_nodes > 0);
        }
        if (num_nodes == 0) {
            return std::min<int>(w_len, OneCharLen(w));
        }
        int mblen = 0;
        for (int i = 0; i < num_nodes; ++i) {
            mblen = std::max<int>(trie_results[i].length, mblen);
        }
        return mblen;
    }

    int PrefixMatcher::PrefixMatch(const std::string &w, bool *found) const {
        return PrefixMatch(w.data(), w.length(), found);
    }

    int PrefixMatcher::PrefixSearch(const char *w, size_t w_len) const {
        if (trie_ == nullptr) {
            return 0;
        }
        constexpr int kResultSize = 64;
        cedar_t::result_pair_type trie_results[kResultSize];
        const int num_nodes = trie_->commonPrefixSearch(w, trie_results, kResultSize, w_len);

        if (num_nodes == 0) {
            return 0;
        }
        int mblen = 0;
        for (int i = 0; i < num_nodes; ++i) {
            mblen = std::max<int>(trie_results[i].length, mblen);
        }
        return mblen;
    }

}  // namespace tbir::runtime
